package io.github.uicdb.et.capability;

import com.mojang.serialization.Codec;
import io.github.uicdb.et.ModElements;
import lombok.RequiredArgsConstructor;
import net.minecraft.core.HolderLookup;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.NbtOps;
import net.neoforged.neoforge.common.util.INBTSerializable;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.UnknownNullability;

import java.util.EnumMap;
import java.util.function.Function;

@RequiredArgsConstructor
public class SimpleElementalHandler implements IElementalHandler, INBTSerializable<CompoundTag> {
    final int maxContainerSize;
    EnumMap<ModElements, Integer> container = new EnumMap<>(ModElements.class);
    private static final Codec<EnumMap<ModElements, Integer>> ENUM_CODEC = Codec.unboundedMap(
            ModElements.CODEC, Codec.INT
    ).xmap(map -> map.isEmpty() ? new EnumMap<>(ModElements.class) : new EnumMap<>(map), Function.identity());

    @Override
    public int drainElemental(int requirement, ModElements type, ElementalAction action) {
        Integer containerOrDefault = container.getOrDefault(type, 0);
        int amount = Math.min(containerOrDefault, requirement);
        if (action.execute()) {
            if (containerOrDefault == amount) {
                container.remove(type);
            } else {
                container.put(type, containerOrDefault - amount);
            }
        }
        onChange(type, containerOrDefault, amount);
        return amount;
    }

    @Override
    public boolean hasElemental(ModElements type) {
        return container.getOrDefault(type, 0) > 0;
    }

    @Override
    public int fillElemental(int requirement, ModElements type, ElementalAction action) {
        Integer has = container.getOrDefault(type, 0);
        int last = maxContainerSize - has;
        int amount = Math.min(last, requirement);
        if (action.execute()) {
            container.put(type, has + amount);
        }
        onChange(type, has, amount);
        return amount;
    }

    @Override
    public int getElementalCapacity(ModElements modElements) {
        return maxContainerSize;
    }

    @Override
    public int getElemental(ModElements modElements) {
        return container.getOrDefault(modElements, 0);
    }

    @Override
    public void removeElemental(ModElements value) {
        onChange(value, container.getOrDefault(value, 0), 0);
        container.remove(value);
    }

    @Override
    public void setElemental(ModElements modElements, int value) {
        onChange(modElements, container.getOrDefault(modElements, 0), value);
        container.put(modElements, value);
    }

    @Override
    public CompoundTag serializeNBT(HolderLookup.@NotNull Provider provider) {
        return (CompoundTag) ENUM_CODEC.encode(container, NbtOps.INSTANCE, new CompoundTag()).getOrThrow();
    }

    @Override
    public void deserializeNBT(HolderLookup.@NotNull Provider provider, @NotNull CompoundTag compoundTag) {
        container = ENUM_CODEC.decode(NbtOps.INSTANCE, compoundTag).getOrThrow().getFirst();
    }

    protected void onChange(ModElements elements, int before, int after) {
    }
}
